import torch
import torch.nn as nn
import torchvision.transforms as transforms
from data import MyDataset
from se_resnext import se_resnext50
import numpy as np
np.set_printoptions(threshold=np.inf)
import os
import sys
sys.path.append('../py2')
from params import data_path

np.set_printoptions(suppress=True)

# 定义是否使用GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 超参数设置
BATCH_SIZE = 128  # 批处理尺寸(batch_size)

# 准备数据集并预处理
transform_train = transforms.Compose([
    # transforms.RandomCrop(32,padding = 4),
    transforms.RandomHorizontalFlip(),  # 图像一半的概率翻转，一半的概率不翻转
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),  # R,G,B每层的归一化用到的均值和方差
])

transform_test = transforms.Compose([
    # transforms.RandomCrop(32,padding = 4),
    transforms.ToTensor(),
    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
])

trainset1 = MyDataset(data_path + '/train1d.txt', transform=transform_train)
# 生成一个个batch进行批训练，组成batch的时候顺序打乱取
trainloader1 = torch.utils.data.DataLoader(trainset1, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)
testset1 = MyDataset(data_path + '/val1d.txt', transform=transform_test)
testloader1 = torch.utils.data.DataLoader(testset1, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)

trainset2 = MyDataset(data_path + '/train2d.txt', transform=transform_train)
# 生成一个个batch进行批训练，组成batch的时候顺序打乱取
trainloader2 = torch.utils.data.DataLoader(trainset2, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)
testset2 = MyDataset(data_path + '/val2d.txt', transform=transform_test)
testloader2 = torch.utils.data.DataLoader(testset2, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)

trainset3 = MyDataset(data_path + '/train3d.txt', transform=transform_train)
# 生成一个个batch进行批训练，组成batch的时候顺序打乱取
trainloader3 = torch.utils.data.DataLoader(trainset3, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)
testset3 = MyDataset(data_path + '/val3d.txt', transform=transform_test)
testloader3 = torch.utils.data.DataLoader(testset3, batch_size=BATCH_SIZE, shuffle=False, num_workers=2)

# 加载保存的模型
net = se_resnext50().to(device)
if torch.cuda.device_count() > 1:
    net=nn.DataParallel(net)
pre = torch.load('model/net_167.pth')
net.load_state_dict(pre)

if __name__ == "__main__":

    # 绘制混淆矩阵
    def confusion_matrix(preds, labels, conf_matrix):
        preds = torch.argmax(preds, 1)
        for p, t in zip(preds, labels):
            conf_matrix[p, t] += 1
        return conf_matrix

    # 测试一下准确率
    conf_matrix = torch.zeros(74, 74)
    print("Waiting Test!")
    with torch.no_grad():
        correct = 0
        total = 0
        for data in testloader3:
            net.eval()
            images, labels = data
            images, labels = images.to(device), labels.to(device)
            outputs = net(images)
            outputs = outputs[2]
            conf_matrix = confusion_matrix(outputs, labels, conf_matrix)
            # 取得分最高的那个类 (outputs.data的索引号)
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum()
            Acc_test = float(10000 * correct / total) / 100
        conf_matrix = conf_matrix.numpy()
        print('测试分类准确率为：%.2f%%' % Acc_test)
        print(conf_matrix)
